import threading
import time
import struct
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import serial.tools.list_ports
import pilink
import pilink_parser

# Common baudrates
default_baudrates = [9600, 19200, 38400, 57600, 115200, 230400]
WORD_SIZE = 4
MAX_RETRY = 5


class BootloaderFlasher:
    def __init__(self, port, baudrate, status_callback=None, progress_callback=None, terminal_callback=None):
        self.port = port
        self.baudrate = baudrate
        self._lock = threading.Lock()
        self.parser = pilink_parser.pilink_serial_parser()
        self.parser.set_msg_callback(self.on_msg)
        self.parser.set_port(port, baudrate)
        self.new_ack = None
        self.new_cmd = None
        self.new_data = None
        self.new_ter = None
        self.status_callback = status_callback
        self.progress_callback = progress_callback
        self.terminal_callback = terminal_callback

    def on_msg(self, port, buf, length):
        if port == pilink_parser.pilink.PILINK_TYPE_BOOT_ACK:
            self.new_ack = pilink_parser.pilink.pilink_boot_ack_decode(buf)
        elif port == pilink_parser.pilink.PILINK_TYPE_BOOT_CMD:
            self.new_cmd = pilink_parser.pilink.pilink_boot_cmd_decode(buf)
        elif port == pilink_parser.pilink.PILINK_TYPE_BOOT_DATA:
            self.new_data = pilink_parser.pilink.pilink_boot_data_decode(buf)
        elif port == pilink_parser.pilink.PILINK_TYPE_TERMINAL:
            rter = pilink_parser.pilink.pilink_terminal_data_decode(buf)
            self.new_ter = rter
            msg = rter.msg[:rter.msg_len].decode('utf-8', errors='ignore')
            if self.terminal_callback:
                self.terminal_callback(msg)

    def _update_status(self, msg):
        if self.status_callback:
            self.status_callback(msg)

    def _update_progress(self, pct):
        if self.progress_callback:
            self.progress_callback(pct)

    def trigger_bootloader(self, timeout=10):
        self._update_status("Waiting for bootloader to start...")
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.parser.send_boot_cmd(pilink_parser.pilink.PILINK_BOOT_CMD_ABORT, 0, 0, 0)
            self.new_ter = None
            time.sleep(0.2)
            self.parser.update()
            if self.new_ter is not None and b"Aborted" in self.new_ter.msg:
                return True
            
        return False

    def erase_flash(self, timeout=10):
        self._update_status("[STEP] try to erase flash")
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.parser.send_boot_cmd(pilink_parser.pilink.PILINK_BOOT_CMD_ERASE, 0, 0, 0)
            self.new_ter = None
            time.sleep(0.2)
            self.parser.update()
            if self.new_ter is not None and b"Eraseing" in self.new_ter.msg:
                break

        self._update_status("[STEP] Erase...")
        deadline = time.time() + 15
        self.new_ter = None
        while time.time() < deadline:
            time.sleep(0.2)
            self.parser.update()
            if self.new_ter is not None and b"Erased" in self.new_ter.msg:
                return True

        return False

    def trigger_upgrade(self, timeout=10):
        self._update_status("[STEP] Start Upgrading...")
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.parser.send_boot_cmd(pilink_parser.pilink.PILINK_BOOT_CMD_STA_UGD, 0, 0, 0)
            self.new_ter = None
            time.sleep(0.2)
            self.parser.update()
            if self.new_ter is not None and b"Upgrading" in self.new_ter.msg:
                return True

        return False

    def end_upgrade(self, timeout=10):
        self._update_status("[STEP] Stop Upgrading...")
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.parser.send_boot_cmd(pilink_parser.pilink.PILINK_BOOT_CMD_END_UGD, 0, 0, 0)
            self.new_ter = None
            time.sleep(0.2)
            self.parser.update()
            if self.new_ter is not None and b"Upgraded" in self.new_ter.msg:
                return True

        return False

    def restart(self, timeout=10):
        with self._lock:
            self._update_status("[STEP] Restart")
            deadline = time.time() + timeout
            while time.time() < deadline:
                self.parser.send_boot_cmd(pilink_parser.pilink.PILINK_BOOT_CMD_RESTART, 0, 0, 0)
                self.new_ter = None
                time.sleep(0.2)
                self.parser.update()
                if self.new_ter is not None and b"Restarting" in self.new_ter.msg:
                    return True
            return False

    def ack_verify(self, addr, expected):
        timeout = 0.5
        deadline = time.time() + timeout
        while time.time() < deadline:
            time.sleep(0.001)
            self.new_ack = None
            self.parser.update()
            if self.new_ack and self.new_ack.addr == addr and self.new_ack.data == expected:
                return True
        return False

    def flash_firmware(self, firmware_path, base_addr):
        # Pre-restart to enter bootloader
        self._update_status("[STEP] Pre-restart board to bootloader")
        if not self.restart():
            self._update_status("[ERROR] Pre-restart failed.")
            return False
        self._update_status("Waiting 3 seconds for bootloader...")
        time.sleep(3)

        # Bootloader detection
        if not self.trigger_bootloader():
            self._update_status("[ERROR] Bootloader not detected. Please reset into bootloader mode.")
            return False
        self._update_status("Bootloader detected.")

        # Erase flash
        if not self.erase_flash():
            self._update_status("[ERROR] Failed to erase flash.")
            return False
        self._update_status("Erasing completed.")

        # Start upgrade session
        if not self.trigger_upgrade():
            self._update_status("[ERROR] Failed to start upgrading.")
            return False
        self._update_status("Upgrading started.")

        # Send firmware data
        self._update_status("[STEP] Sending firmware data")
        with open(firmware_path, 'rb') as f:
            data = f.read()
        if len(data) % WORD_SIZE:
            data += b'\xFF' * (WORD_SIZE - len(data) % WORD_SIZE)
        data_len = len(data)
        for offset in range(0, data_len, WORD_SIZE):
            word = struct.unpack('<I', data[offset:offset+4])[0]
            addr = base_addr + offset
            pct = offset * 100 / data_len
            self._update_progress(pct)
            for attempt in range(1, MAX_RETRY+1):
                self.parser.send_boot_data(addr, word)
                if self.ack_verify(addr, word):
                    break
                self._update_status(f"[WARN] Retry {attempt} for address 0x{addr:08X}")
            else:
                self._update_status(f"[ERROR] Failed to program at address 0x{addr:08X}")
                return False

        # Finalize and restart
        self._update_status("[STEP] Finalizing upgrade")
        if not self.end_upgrade():
            self._update_status("[ERROR] Failed to stop upgrading.")
            return False
        self._update_status("Upgrading successfully.")
        self._update_status("Restarting target system")
        self.restart()
        return True

class FlashGUI:
    def __init__(self, root):
        self.root = root
        root.title("Bootloader Flash Tool")

        main_frame = ttk.Frame(root, padding=(15, 15, 15, 15))
        main_frame.pack(fill='both', expand=True)

        # Serial Port
        ttk.Label(main_frame, text="Serial Port:").grid(row=0, column=0, sticky='e', padx=5, pady=5)
        self.port_cb = ttk.Combobox(main_frame, values=self.scan_ports())
        self.port_cb.grid(row=0, column=1, columnspan=2, sticky='we', padx=5)

        # Baudrate
        ttk.Label(main_frame, text="Baudrate:").grid(row=1, column=0, sticky='e', padx=5, pady=5)
        self.baud_cb = ttk.Combobox(main_frame, values=default_baudrates)
        self.baud_cb.set(115200)
        self.baud_cb.grid(row=1, column=1, columnspan=2, sticky='we', padx=5)

        # Firmware File
        ttk.Label(main_frame, text="Firmware File:").grid(row=2, column=0, sticky='e', padx=5, pady=5)
        self.file_var = tk.StringVar()
        ttk.Entry(main_frame, textvariable=self.file_var).grid(row=2, column=1, sticky='we', padx=5)
        ttk.Button(main_frame, text="Browse", command=self.browse_file).grid(row=2, column=2, padx=5)

        # Start Address
        ttk.Label(main_frame, text="Start Address:").grid(row=3, column=0, sticky='e', padx=5, pady=5)
        self.base_var = tk.StringVar(value="0x08008000")
        ttk.Entry(main_frame, textvariable=self.base_var).grid(row=3, column=1, columnspan=2, sticky='we', padx=5)

        # Buttons
        btn_frame = ttk.Frame(main_frame)
        btn_frame.grid(row=4, column=0, columnspan=3, pady=10)
        self.start_btn = ttk.Button(btn_frame, text="Start Flash", command=self.start_flash)
        self.start_btn.pack(side='left', padx=10)
        self.restart_btn = ttk.Button(btn_frame, text="Restart Board", command=self.start_restart)
        self.restart_btn.pack(side='left', padx=10)

        # Status label and Clear Logs
        self.status_var = tk.StringVar(value="Idle")
        ttk.Label(main_frame, textvariable=self.status_var).grid(row=5, column=0, columnspan=2, sticky='w', pady=(5,2))
        ttk.Button(main_frame, text="Clear Logs", command=self.clear_logs).grid(row=5, column=2, sticky='e', padx=5)

        # Progress bar
        self.progress = ttk.Progressbar(main_frame, orient='horizontal', mode='determinate')
        self.progress.grid(row=6, column=0, columnspan=3, sticky='we', pady=(2,10))

        # Logs Title
        ttk.Label(main_frame, text="--- Logs ---", font=(None, 11, 'bold')).grid(row=7, column=0, columnspan=3, pady=(0,5))

        # Status Log (top)
        ttk.Label(main_frame, text="Status Log:").grid(row=8, column=0, columnspan=3, sticky='w', padx=5)
        self.status_text = tk.Text(main_frame, height=6, wrap='none')
        self.status_text.grid(row=9, column=0, columnspan=3, sticky='nsew', padx=5)
        st_sc = ttk.Scrollbar(main_frame, orient='vertical', command=self.status_text.yview)
        st_sc.grid(row=9, column=3, sticky='ns', pady=5)
        self.status_text.config(yscrollcommand=st_sc.set)

        # Terminal Log (bottom)
        ttk.Label(main_frame, text="Terminal Log:").grid(row=10, column=0, columnspan=3, sticky='w', padx=5, pady=(10,0))
        self.terminal_text = tk.Text(main_frame, height=6, wrap='none')
        self.terminal_text.grid(row=11, column=0, columnspan=3, sticky='nsew', padx=5, pady=(0,10))
        ter_sc = ttk.Scrollbar(main_frame, orient='vertical', command=self.terminal_text.yview)
        ter_sc.grid(row=11, column=3, sticky='ns', pady=5)
        self.terminal_text.config(yscrollcommand=ter_sc.set)

        # Configure stretching
        main_frame.columnconfigure(1, weight=1)
        main_frame.rowconfigure(9, weight=1)
        main_frame.rowconfigure(11, weight=1)

        self.flasher = None

    def scan_ports(self):
        return [p.device for p in serial.tools.list_ports.comports()]

    def browse_file(self):
        p = filedialog.askopenfilename(filetypes=[("BIN", "*.bin"), ("All", "*")])
        if p:
            self.file_var.set(p)

    def ensure_flasher(self):
        port = self.port_cb.get()
        baud = int(self.baud_cb.get())
        if not self.flasher or self.flasher.port != port or self.flasher.baudrate != baud:
            self.flasher = BootloaderFlasher(
                port, baud,
                status_callback=self.log_status,
                progress_callback=lambda v: self.progress.step(v - self.progress['value']),
                terminal_callback=self.log_terminal
            )
        return self.flasher

    def log_status(self, msg):
        self.status_var.set(msg)
        self.status_text.insert(tk.END, msg + "\n")
        self.status_text.see(tk.END)

    def log_terminal(self, msg):
        self.terminal_text.insert(tk.END, msg)
        self.terminal_text.see(tk.END)

    def clear_logs(self):
        self.status_text.delete('1.0', tk.END)
        self.terminal_text.delete('1.0', tk.END)

    def start_flash(self):
        port = self.port_cb.get()
        fw = self.file_var.get()
        base = int(self.base_var.get(), 0)
        if not port or not fw:
            messagebox.showerror("Error", "Select port and firmware.")
            return

        self.start_btn.config(state='disabled')
        self.log_status("Starting flash...")
        self.progress['value'] = 0

        fl = self.ensure_flasher()
        threading.Thread(target=lambda: self._flash_worker(fl, fw, base), daemon=True).start()

    def _flash_worker(self, fl, fw, base):
        fl.restart()
        time.sleep(3)
        ok = fl.flash_firmware(fw, base)
        if ok:
            messagebox.showinfo("Success", "Firmware updated.")
        else:
            messagebox.showerror("Fail", "Update failed.")
        self.start_btn.config(state='normal')

    def start_restart(self):
        port = self.port_cb.get()
        if not port:
            messagebox.showerror("Error", "Select port.")
            return

        self.restart_btn.config(state='disabled')
        self.log_status("Restarting board...")

        fl = self.ensure_flasher()
        threading.Thread(target=lambda: self._restart_worker(fl), daemon=True).start()

    def _restart_worker(self, fl):
        ok = fl.restart()
        if ok:
            messagebox.showinfo("Success", "Board restarting.")
        else:
            messagebox.showerror("Fail", "Failed to restart board.")
        self.restart_btn.config(state='normal')


if __name__=='__main__':
    root = tk.Tk()
    FlashGUI(root)
    root.mainloop()
